import json
import random
import secrets
import string
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
import base64
import os
from app import app


def generate_key(as_base64: bool = True) -> str:
    """
    生成一个适用于AES加密的安全随机密钥
    
    Args:
        as_base64 (bool): 是否返回Base64编码的密钥字符串。默认为True
        
    Returns:
        str: 如果as_base64为True，返回Base64编码的密钥字符串；
             如果as_base64为False，返回十六进制编码的密钥字符串
    """
    # 生成32字节(256位)的随机密钥
    key = os.urandom(32)

    if as_base64:
        # 返回Base64编码的密钥
        return base64.b64encode(key).decode('utf-8')
    else:
        # 返回十六进制编码的密钥
        return key.hex()

def generate_key_2(length=16):
    """生成指定长度的随机字符串（仅包含字母和数字）"""
    alphabet = string.ascii_letters + string.digits
    return ''.join(secrets.choice(alphabet) for i in range(length))

'''
Python 后端对返回前端的数据进行加密的代码
'''


def encrypt_data(data: str, key: str) -> str:
    # 将密钥转换为 bytes
    key_bytes = key.encode('utf-8').ljust(32, b'\0')[:32]

    # 生成随机初始化向量
    iv = os.urandom(16)

    # 创建加密器
    cipher = Cipher(
        algorithms.AES(key_bytes),
        modes.CBC(iv),
        backend=default_backend()
    )
    encryptor = cipher.encryptor()

    # 填充数据
    padder = padding.PKCS7(128).padder()
    padded_data = padder.update(data.encode('utf-8')) + padder.finalize()

    # 加密数据
    ciphertext = encryptor.update(padded_data) + encryptor.finalize()

    # 组合 IV 和密文并进行 Base64 编码
    combined = iv + ciphertext
    return base64.b64encode(combined).decode('utf-8')


def encrypt_data_2(text, key):
    """
    使用AES-256-GCM算法加密字符串

    Args:
        text (str): 要加密的明文字符串
        key (str): 不定长密钥
    Returns:
        str: JSON格式的加密结果，包含salt、nonce和加密数据
    """
    if not text or not key:
        raise ValueError("文本和密钥都不能为空")

    # 生成随机salt和nonce
    salt = os.urandom(16)  # 128位salt
    nonce = os.urandom(12)  # 96位nonce (GCM推荐)

    # 使用PBKDF2从密钥派生256位密钥
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,  # 256位密钥
        salt=salt,
        iterations=100000,  # 10万次迭代
        backend=default_backend()
    )
    derived_key = kdf.derive(key.encode('utf-8'))

    # 创建AES-GCM加密器
    cipher = Cipher(
        algorithms.AES(derived_key),
        modes.GCM(nonce),
        backend=default_backend()
    )
    encryptor = cipher.encryptor()

    # 加密数据
    plaintext_bytes = text.encode('utf-8')
    ciphertext = encryptor.update(plaintext_bytes) + encryptor.finalize()

    # 获取认证标签
    auth_tag = encryptor.tag

    # 构造结果
    result = {
        'salt': base64.b64encode(salt).decode('utf-8'),
        'nonce': base64.b64encode(nonce).decode('utf-8'),
        'ciphertext': base64.b64encode(ciphertext).decode('utf-8'),
        'tag': base64.b64encode(auth_tag).decode('utf-8')
    }

    return json.dumps(result)

def simple_encrypt(data: str) -> str:
    return encrypt_data(data, app.config['SECRET_KEY'])


def complex_encrypt(data: str) -> str:
    key_len = random.randint(8, 32)
    key = generate_key_2(key_len)
    data = encrypt_data_2(data, key)
    offset = random.randint(0, key_len + len(data))
    encrypted_data = (str(len(str(key_len))) + str(key_len) +
                      str(len(str(offset))) + str(offset) +
                      data[0:offset] + key + data[offset:])
    return encrypted_data
